
/*------------------------------------------

A suffix tree implementation for PAT-Array

------------------------------------------*/
#ifndef __SUFFIX_TREE_H
#define __SUFFIX_TREE_H

#include <math.h>
#include "sentence.h"
#include "AVL.h"

template <class KeyType>
struct SuffixTreeNode;

// used to collate context information
template <class KeyType>
struct WordContext
{
	wchar_t *str;
	int length;

	SuffixTreeNode<KeyType> * treeNode;

	WordContext()
	{
		length = 0;
		str=0;
		treeNode=0;
	}

	~WordContext()
	{
		delete [] str;
	}
};

template <class KeyType>
struct SuffixTreeNode
{
	KeyType state;
	AVL<KeyType, SuffixTreeNode<KeyType> *> children;
	AVL<KeyType, KeyType*> tails;

	int counter;
	int endWithNullCounter;
	int tailsSize;
	int childrenSize;

	SuffixTreeNode()
	{
		//state = 0;
		counter = 0;
		childrenSize=tailsSize=endWithNullCounter=0;
	}
	
	~SuffixTreeNode()
	{	
		
	}

	void free()
	{
	}
	
	SuffixTreeNode* findChildNode(KeyType key)
	{
		SuffixTreeNode<KeyType> * item = 0;
		children.Find(key,item);
		return item;
	}

	SuffixTreeNode* addChildNode(KeyType state)
	{
		SuffixTreeNode<KeyType> * node = new SuffixTreeNode<KeyType>();
		node->state = state;
		
		children.Insert(state, node);
		childrenSize++;

		return node;
	}

	KeyType* findTail(KeyType key)
	{
		KeyType *item = 0;
		tails.Find(key, item);
		return item;
	}
	void addStringTail(KeyType * tail)
	{
		tails.Insert(*tail,tail);
		tailsSize++;
	}
	void removeTail(KeyType tail)
	{
		tails.Remove(tail);
		tailsSize--;
	}

	double entropy()
	{
		double result=0.0;
		double totalTimes = this->counter;

		// children
		KeyType key;
		SuffixTreeNode<KeyType> * item;
		AVL<KeyType, SuffixTreeNode<KeyType> *>::Iterator<KeyType, SuffixTreeNode<KeyType> *> iter(&children);
		bool more = iter.GetFirst(key,item);
		while(more)
		{
			int times = item->counter;
			result += times*log(times/totalTimes);
			more = iter.GetNext(key,item);
		}

		//tail string
		int num = tailsSize;
		if(num > 0)
		{
			result += num * 1 * log(1/totalTimes);
		}
		//some times this node is end of string.
		int times = this->endWithNullCounter;
		if(times > 0)
		{
			result += times*log(times/totalTimes);
		}

		return abs(result);
	}
};


template <class KeyType>
class SuffixTree
{
public:
	//define wheter it is usted to represent left pat-array or right pat-array.
	enum TREE_TYPE { LEFT_SUFFIX_TREE=0xBABE001, RIGHT_SUFFIX_TREE=0xBABE002 };

private:
	int minMumWordLength;
	SuffixTreeNode<KeyType> *root;
	TREE_TYPE type;	

public:

	SuffixTree(TREE_TYPE setType ){
		minMumWordLength = 2;
		root = new SuffixTreeNode<KeyType>();
		type = setType;
	}

	~SuffixTree()
	{
		root->free();
		delete root;
	}
	
	void insert(KeyType* suffix);
	void collectWordContext(WordContext<KeyType>* & context,int& size);
	void travel();
	
private:
	// insert at a tree node
	void insert(SuffixTreeNode<KeyType> * cur, KeyType *suffix);
	void appendAsChildOrStringTail(SuffixTreeNode<KeyType> * cur, KeyType *suffix);
	void travel(SuffixTreeNode<KeyType> *node,wchar_t *prefix, int prefixLength);

	/*
	 * node :		current tree root.
	 * curdepth:	current root's parent depth in tree.
	 * prefix:		the prefix string of current string(exclude the state stored in itself)
	 * list:		used to remember all data.
	 * index:		the index of list where to store the data.
	 */
	void collectWordContext(SuffixTreeNode<KeyType> *node, int curdepth, wchar_t* prefix, int prefixLength, WordContext<KeyType>* list,int& index);
	void appendToPrefix(wchar_t * prefix, KeyType state, wchar_t *newResult, int length);

	int countNodes(SuffixTreeNode<KeyType> *root,int curdepth);
	int countWordInserted(SuffixTreeNode<KeyType> *root);
};	

template<class KeyType>
void SuffixTree<KeyType>::insert(KeyType* suffix)
{
	appendAsChildOrStringTail(root, suffix);
}


template<class KeyType>
void SuffixTree<KeyType>::appendAsChildOrStringTail(SuffixTreeNode<KeyType> * cur, KeyType * suffix)
{

	//whether the Key is in Children nodes.
	SuffixTreeNode<KeyType> * node=0;
	if( (node = cur->findChildNode(*suffix)) != 0)
	{
		insert(node,suffix);
		return;
	}

	KeyType * tail = 0;
	// if the Key is in Tails.
	if(( tail = cur->findTail(*suffix)) != 0)
	{
		cur->removeTail(*tail);
		SuffixTreeNode<KeyType> *newNode = cur->addChildNode(*suffix);
		
		insert(newNode, tail);
		insert(newNode, suffix);
	}
	// then add as a new tail appended to current node.
	else
	{
		cur->addStringTail(suffix);
	}
}

template<class KeyType>
void SuffixTree<KeyType>::collectWordContext(WordContext<KeyType> *& list, int& size)
{

	// we'll only collect the context of word whose depth(length) is larger than 2.
	int listSize = countNodes(root,-1);
	list = new WordContext<KeyType>[listSize];
	
	size = 0;		// in this loop it is sure that size will increment to listSize.
	AVL<KeyType, SuffixTreeNode<KeyType> *>::Iterator<KeyType, SuffixTreeNode<KeyType> *> iter(&(root->children));
	KeyType key;
	SuffixTreeNode<KeyType> *item;
	bool more = iter.GetFirst(key,item);
	while(more)
	{
		//printf("Root Children....%d\n",size);
		collectWordContext(item,0,0,0,list,size);
		more = iter.GetNext(key,item);
	}
}

//
//template<>
//void SuffixTree<SegmentFragment>::collectWordContext(SuffixTreeNode<SegmentFragment> *node, int parent_depth, wchar_t* prefix, WordContext<SegmentFragment>* list,int& index)
//{
//	//int depth = parent_depth+1;
//	//
//	//if(depth < minMumWordLength)
//	//{
//	//	wchar_t * temp = new wchar_t[depth];
//	//	appendToPrefix(prefix,node->state,temp,parent_depth);
//
//	//	//travel to all node recursively
//	//	for( int i = 0; i < node->children_size; ++i)
//	//	{
//	//		collectWordContext(&node->nodeList[i],depth, temp,list,index);
//	//	}
//	//	delete [] temp;
//	//}
//	//else
//	//{
//	//	WordContext * ctx = &list[index++];
//
//	//	//initial length
//	//	ctx->length = depth;
//	//	
//	//	//initial str
//	//	ctx->str = new wchar_t[depth];
//	//	
//	//	appendToPrefix(prefix,node->state,ctx->str, parent_depth);
//
//	//	//remember current node reference
//	//	ctx->treeNode = node;
//
//	//	//travel to all node recursively
//	//	for( int i = 0; i < node->children_size; ++i)
//	//	{
//	//		collectWordContext(&node->nodeList[i],depth,ctx->str,list,index);
//	//	}
//	//}
//}
//
//

template <class KeyType>
int SuffixTree<KeyType>::countNodes(SuffixTreeNode<KeyType> *root,int parent_depth)
{
	int counter = 0;
	int curdepth = parent_depth+1;

	if( curdepth >= minMumWordLength){
		counter = 1;
	}
	
	AVL<KeyType, SuffixTreeNode<KeyType> *>::Iterator<KeyType, SuffixTreeNode<KeyType> *> iter(&root->children);
	KeyType key;
	SuffixTreeNode<KeyType> *item;
	bool more = iter.GetFirst(key,item);
	while(more)
	{
		counter +=  countNodes(item,curdepth);
		more = iter.GetNext(key,item);
	}
	
	return counter;                                                                                                                                                                                                                                                                                           ;
}

template <class KeyType>
int SuffixTree<KeyType>::countWordInserted(SuffixTreeNode<KeyType> *root)
{
	int count=root->tailsSize + root->endWithNullCounter;
	
	AVL<KeyType, SuffixTreeNode<KeyType> *>::Iterator<KeyType, SuffixTreeNode<KeyType> *> iter(&root->children);
	KeyType key;
	SuffixTreeNode<KeyType> *item;
	bool more = iter.GetFirst(key,item);
	while(more)
	{
		count += countWordInserted(item);
		more = iter.GetNext(key,item);
	}
	return count;

}


//------------------------------------------------------


#endif